import os
import json
from py2neo import Graph, Node

import pandas as pd

df = pd.read_csv('medical_data3.csv', encoding='gbk')

# 实体——所有症状
symptoms = []
for each in df['症状']:
    symptoms.extend(each.split(','))
symptoms = set(symptoms)

# 实体——疾病种类（上位）
types = []
for each in df['疾病种类']:
    types.extend(each.split(','))
types = set(types)

# 实体——药治
drugs = []
for each in df['药治']:
    try:
        drugs.extend(each.split(','))
    except:
        pass
drugs = set(drugs)

# 实体——并发症
complication = []
for each in df['并发症']:
    try:
        complication.extend(each.split(','))
    except:
        pass
complication = set(complication)

# 疾病字典信息
disease_infos = []  # 疾病信息
for idx, row in df.iterrows():
    disease_infos.append(dict(row))
# dict(row).keys()
    
def deduplicate(rels_old):
    '''关系去重函数'''
    rels_new = []
    for each in rels_old:
        if each not in rels_new:
            rels_new.append(each)
    return rels_new

# 关系：疾病——病因
rels_reason = []
for idx, row in df.iterrows():
    try:
        for each in row['病因'].split(','):
            rels_reason.append([row['疾病名称'], each])
    except:
        pass
rels_reason = deduplicate(rels_reason)
# rels_reason

# 关系：疾病——症状
rels_symptom = []
for idx, row in df.iterrows():
    for each in row['症状'].split(','):
        rels_symptom.append([row['疾病名称'], each])
rels_symptom = deduplicate(rels_symptom)
# rels_symptom

# 关系：疾病——疾病（并发症）
rels_acompany = []
for idx, row in df.iterrows():
    for each in row['并发症'].split(','):
        rels_acompany.append([row['疾病名称'], each])
rels_acompany = deduplicate(rels_acompany)

# 关系：疾病——药物（药治）
rels_drug = []
for idx, row in df.iterrows():
    try:
        for each in row['药治'].split(','):
            rels_drug.append([row['疾病名称'], each])
    except:
        pass
rels_recommanddrug = deduplicate(rels_drug)

# 关系：疾病——疾病种类（上下位）
rels_category = []  # 关系：疾病-种类
rels_department = []  # 关系：小上位-大上位（例如：呼吸内科属于内科的一种）
for idx, row in df.iterrows():
    if len(row['疾病种类'].split(',')) == 1:
        rels_category.append([row['疾病名称'], row['疾病种类']])
    else:
        big = row['疾病种类'].split(',')[0] # 大上位
        small = row['疾病种类'].split(',')[1] # 小上位
        rels_category.append([row['疾病名称'], small])
        rels_department.append([small, big])
rels_category = deduplicate(rels_category)
rels_department = deduplicate(rels_department)
# rels_category
# rels_department


# 连接知识图谱 开始操作
g = Graph('http://localhost:7474', auth=('neo4j', '123456'))
# 创建疾病实体
count = 0
for disease_dict in disease_infos:
    try:
        node = Node("Disease",
                    name=disease_dict['疾病名称'],
                    cause=disease_dict['病因'],
                    disease_type=disease_dict['疾病种类'])
        g.create(node)
        count += 1
        print('创建疾病实体：', disease_dict['疾病名称'])
    except:
        pass
print('共创建 {} 个疾病实体'.format(count))

# 创建并发症实体
for each in complication:
    node = Node('Complication', name=each)
    g.create(node)
    print('创建实体 {}'.format(each))

# 创建病因实体
for each in rels_reason:
    node = Node('Reason', name=each)
    g.create(node)
    print('创建实体 {}'.format(each))

# 创建药物实体
for each in drugs:
    node = Node('Drug', name=each)
    g.create(node)
    print('创建实体 {}'.format(each))

# 创建疾病种类实体
for each in types:
    node = Node('Type', name=each)
    g.create(node)
    print('创建实体 {}'.format(each))

# 创建症状实体
for each in symptoms:
    node = Node('Symptom', name=each)
    g.create(node)
    print('创建实体 {}'.format(each))

def create_relationship(start_node, end_node, edges, rel_type, rel_name):
    '''创建知识图谱关系函数'''
    for edge in edges:
        p = edge[0]
        q = edge[1]
        # 创建关系的 Cypher 语句
        query = "match(p:%s),(q:%s) where p.name='%s' and q.name='%s' create (p)-[rel:%s{name:'%s'}]->(q)" % (start_node, end_node, p, q, rel_type, rel_name)
        try:
            g.run(query) # 运行 Cypher 语句
            print('创建关系 {}-{}->{}'.format(p, rel_type, q))
        except Exception as e:
            print(e)

# 构建实体对关系
create_relationship('Disease', 'Reason', rels_reason, 'cause', '病因')
create_relationship('Disease', 'Drug', rels_drug, 'drug', '药治')
create_relationship('Type', 'Type', rels_department, 'belongs_to', '属于')
create_relationship('Disease', 'Symptom', rels_symptom, 'has_symptom', '症状')
create_relationship('Disease', 'Complication', rels_acompany, 'acompany_with', '并发症')
create_relationship('Disease', 'Type', rels_category, 'belongs_to', '疾病种类')
# 删除所有节点
# cypher = 'MATCH (n) DETACH DELETE n'
# g.run(cypher)